﻿#include <iostream>
#include <algorithm>
#include <strsafe.h>
#include <stdio.h>
#include "CLogUtils.h"

#include <tlhelp32.h>
#include <tchar.h>
#include <string>

#ifdef UNICODE
#define _tstring		std::wstring
#else
#define _tstring		std::string
#endif

static DWORD DuplicateWinloginToken(DWORD dwSessionId, DWORD dwDesiredAccess, PHANDLE phToken) {
    DWORD dwErr;
    PRIVILEGE_SET ps;

    ps.PrivilegeCount = 1;
    ps.Control = PRIVILEGE_SET_ALL_NECESSARY;

    if (LookupPrivilegeValue(NULL, SE_TCB_NAME, &ps.Privilege[0].Luid)) {
        HANDLE hSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if (INVALID_HANDLE_VALUE != hSnapshot) {
            BOOL bCont, bFound = FALSE;
            PROCESSENTRY32 pe;

            pe.dwSize = sizeof(pe);
            dwErr = ERROR_NOT_FOUND;

            for (bCont = Process32First(hSnapshot, &pe); bCont; bCont = Process32Next(hSnapshot, &pe)) {
                HANDLE hProcess;

                if (0 != _tcsicmp(pe.szExeFile, TEXT("winlogon.exe"))) {
                    continue;
                }

                hProcess = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, FALSE, pe.th32ProcessID);
                if (hProcess) {
                    HANDLE hToken;
                    DWORD dwRetLen, sid;

                    if (OpenProcessToken(hProcess, TOKEN_QUERY | TOKEN_DUPLICATE, &hToken)) {
                        BOOL fTcb;

                        if (PrivilegeCheck(hToken, &ps, &fTcb) && fTcb) {
                            if (GetTokenInformation(hToken, TokenSessionId, &sid, sizeof(sid), &dwRetLen) && sid == dwSessionId) {
                                bFound = TRUE;
                                if (DuplicateTokenEx(hToken, dwDesiredAccess, NULL, SecurityImpersonation, TokenImpersonation, phToken)) {
                                    dwErr = ERROR_SUCCESS;
                                }
                                else {
                                    dwErr = GetLastError();
                                }
                            }
                        }
                        CloseHandle(hToken);
                    }
                    CloseHandle(hProcess);
                }

                if (bFound) break;
            }

            CloseHandle(hSnapshot);
        }
        else {
            dwErr = GetLastError();
        }
    }
    else {
        dwErr = GetLastError();
    }


    return dwErr;
}

static DWORD CreateUIAccessToken(PHANDLE phToken) {
    DWORD dwErr;
    HANDLE hTokenSelf;

    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY | TOKEN_DUPLICATE, &hTokenSelf)) {
        DWORD dwSessionId, dwRetLen;

        if (GetTokenInformation(hTokenSelf, TokenSessionId, &dwSessionId, sizeof(dwSessionId), &dwRetLen)) {
            HANDLE hTokenSystem;

            dwErr = DuplicateWinloginToken(dwSessionId, TOKEN_IMPERSONATE, &hTokenSystem);
            if (ERROR_SUCCESS == dwErr) {
                if (SetThreadToken(NULL, hTokenSystem)) {
                    if (DuplicateTokenEx(hTokenSelf, TOKEN_QUERY | TOKEN_DUPLICATE | TOKEN_ASSIGN_PRIMARY | TOKEN_ADJUST_DEFAULT, NULL, SecurityAnonymous, TokenPrimary, phToken)) {
                        BOOL bUIAccess = TRUE;

                        if (!SetTokenInformation(*phToken, TokenUIAccess, &bUIAccess, sizeof(bUIAccess))) {
                            dwErr = GetLastError();
                            CloseHandle(*phToken);
                        }
                    }
                    else {
                        dwErr = GetLastError();
                    }
                    RevertToSelf();
                }
                else {
                    dwErr = GetLastError();
                }
                CloseHandle(hTokenSystem);
            }
        }
        else {
            dwErr = GetLastError();
        }

        CloseHandle(hTokenSelf);
    }
    else {
        dwErr = GetLastError();
    }

    return dwErr;
}

static _tstring GetModuleDir()
{
    TCHAR szCurPath[MAX_PATH] = { 0 };
    ::GetModuleFileName(NULL, szCurPath, _countof(szCurPath));
    _tstring strDir = szCurPath;
    strDir = strDir.substr(0, strDir.find_last_of(_T('\\')));
    return strDir;
}

static _tstring GetModulePath()
{
    TCHAR szCurPath[MAX_PATH] = { 0 };
    ::GetModuleFileName(NULL, szCurPath, _countof(szCurPath));
    return szCurPath;
}

static BOOL CheckForUIAccess(DWORD* pdwErr, DWORD* pfUIAccess) {
    BOOL result = FALSE;
    HANDLE hToken;

    if (OpenProcessToken(GetCurrentProcess(), TOKEN_QUERY, &hToken)) {
        DWORD dwRetLen;

        if (GetTokenInformation(hToken, TokenUIAccess, pfUIAccess, sizeof(*pfUIAccess), &dwRetLen)) {
            result = TRUE;
        }
        else {
            *pdwErr = GetLastError();
        }
        CloseHandle(hToken);
    }
    else {
        *pdwErr = GetLastError();
    }

    return result;
}

DWORD PrepareForUIAccess() {
    DWORD dwErr;
    HANDLE hTokenUIAccess;
    DWORD fUIAccess;

    if (CheckForUIAccess(&dwErr, &fUIAccess)) {
        if (fUIAccess) {
            dwErr = ERROR_SUCCESS;
        }
        else {
            dwErr = CreateUIAccessToken(&hTokenUIAccess);
            if (ERROR_SUCCESS == dwErr) {
                STARTUPINFO si;
                PROCESS_INFORMATION pi;

                ::GetStartupInfo(&si);
                _tstring exePath = GetModulePath();
                _tstring lpCommandLine = ::GetCommandLine();
                lpCommandLine += _T(" Launch");
                if (::CreateProcessAsUser(hTokenUIAccess, exePath.c_str(), (LPTSTR)lpCommandLine.c_str(), NULL, NULL, FALSE, 0, NULL, NULL, &si, &pi)) {
                    CloseHandle(pi.hProcess), CloseHandle(pi.hThread);
                    ExitProcess(0);
                }
                else {
                    dwErr = GetLastError();
                }

                CloseHandle(hTokenUIAccess);
            }
        }
    }

    return dwErr;
}

void EnumDesktopWnd(HDESK hDesktop)
{
    DWORD dwIndex = 0;
    ::EnumDesktopWindows(hDesktop, [](HWND hWnd, LPARAM lParam) -> BOOL {
        LPDWORD lpIndex = (LPDWORD)lParam;

        TCHAR szWndName[MAX_PATH] = { 0 };
        TCHAR szClassName[MAX_PATH] = { 0 };
        RECT rtWnd = { 0 };
        ::GetWindowText(hWnd, szWndName, _countof(szWndName));
        ::GetClassName(hWnd, szClassName, _countof(szClassName));
        ::GetWindowRect(hWnd, &rtWnd);

        if (!::IsWindowVisible(hWnd))
        {
            return TRUE;
        }

        // 过滤系统窗口与隐藏窗口
        DWORD dwFilterExStyle = WS_EX_TOOLWINDOW | WS_EX_NOACTIVATE;
        if ((::GetWindowLongPtr(hWnd, GWL_EXSTYLE) & dwFilterExStyle))
        {
            return TRUE;
        }

        LPCTSTR ClassNameFilter[] = {
            _T("Windows.Internal.Shell.TabProxyWindow"),
            //_T("Windows.UI.Core.CoreWindow"),
            _T("CASCADIA_HOSTING_WINDOW_CLASS")
        
        };

        for (int i = 0; i < _countof(ClassNameFilter); i++)
        {
            if (0 == _tcsicmp(ClassNameFilter[i], szClassName))
            {
                return TRUE;
            }
        }

        if (::IsIconic(hWnd))
        {
            return TRUE;
        }

        HWND hPointWnd = ::WindowFromPoint(POINT{ rtWnd.left + (rtWnd.right - rtWnd.left), rtWnd.top + (rtWnd.bottom - rtWnd.top)});
        if (hPointWnd == hWnd)
        {
            CLogUtils::LOG_PRINT(_T("%04d HWND: %08X hPointWnd: %08X name: %s class: %s"), (*lpIndex)++, hWnd, hPointWnd, szWndName, szClassName);
        }
        return TRUE;

        }, (LPARAM)&dwIndex);
}

int APIENTRY _tWinMain(
    _In_ HINSTANCE hInstance,
    _In_opt_ HINSTANCE hPrevInstance,
    _In_ LPTSTR lpCmdLine,
    _In_ int nCmdShow
)
{
    UNREFERENCED_PARAMETER(hPrevInstance);
    UNREFERENCED_PARAMETER(lpCmdLine);
    UNREFERENCED_PARAMETER(nCmdShow);

    if (lpCmdLine[0] == _T('\0'))
    {
        PrepareForUIAccess();
        return 0;
    }

    CLogUtils::GetInstance().Initialize(true, _T("AYASpace"), _T(""), _T("AYASpace"), CLogUtils::eLogOutputType::eLogOutputNone);

    HDESK hDesktop = ::OpenDesktop(_T("default"), DF_ALLOWOTHERACCOUNTHOOK, FALSE, DESKTOP_ENUMERATE);

    if (NULL != hDesktop)
    {
        while (true)
        {
            /*{

                GUITHREADINFO gti = { sizeof(GUITHREADINFO) };
                GetGUIThreadInfo(0, &gti);

                TCHAR szWndName[MAX_PATH] = { 0 };
                TCHAR szClassName[MAX_PATH] = { 0 };

                ::GetWindowText(gti.hwndActive, szWndName, _countof(szWndName));
                ::GetClassName(gti.hwndActive, szClassName, _countof(szClassName));

                CLogUtils::LOG_PRINT(_T("======================Active==========================="));

                CLogUtils::LOG_PRINT(_T("Active: %08X name: %s class: %s"), gti.hwndActive, szWndName, szClassName);
            }*/

            /*{
                HWND hCurForegroundWnd = ::GetForegroundWindow();

                TCHAR szWndName[MAX_PATH] = { 0 };
                TCHAR szClassName[MAX_PATH] = { 0 };

                ::GetWindowText(hCurForegroundWnd, szWndName, _countof(szWndName));
                ::GetClassName(hCurForegroundWnd, szClassName, _countof(szClassName));

                CLogUtils::LOG_PRINT(_T("======================Foreground==========================="));
                CLogUtils::LOG_PRINT(_T("Foreground: %08X name: %s class: %s"), hCurForegroundWnd, szWndName, szClassName);
            }*/

            {
                HWND hWnd = ::WindowFromPoint(POINT{ 960, 540 });
                TCHAR szWndName[MAX_PATH] = { 0 };
                TCHAR szClassName[MAX_PATH] = { 0 };

                ::GetWindowText(hWnd, szWndName, _countof(szWndName));
                ::GetClassName(hWnd, szClassName, _countof(szClassName));

                CLogUtils::LOG_PRINT(_T("======================WindowFromPoint==========================="));
                CLogUtils::LOG_PRINT(_T("WindowFromPoint: %08X name: %s class: %s"), hWnd, szWndName, szClassName);
            }

            /*{
                HWND hWnd = ::WindowFromPhysicalPoint(POINT{ 960, 540 });
                TCHAR szWndName[MAX_PATH] = { 0 };
                TCHAR szClassName[MAX_PATH] = { 0 };

                ::GetWindowText(hWnd, szWndName, _countof(szWndName));
                ::GetClassName(hWnd, szClassName, _countof(szClassName));

                CLogUtils::LOG_PRINT(_T("======================WindowFromPhysicalPoint==========================="));
                CLogUtils::LOG_PRINT(_T("WindowFromPoint: %08X name: %s class: %s"), hWnd, szWndName, szClassName);
            }*/

            //CLogUtils::LOG_PRINT(_T("========================EnumDesktopWnd========================="));
            //EnumDesktopWnd(hDesktop);
            ::Sleep(1000);
        }
    }

    return 0;
}
